{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[View the runnable example on GitHub](https://github.com/intel-analytics/BigDL/tree/main/python/nano/tutorial/notebook/inference/pytorch/accelerate_pytorch_inference_openvino.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Accelerate PyTorch Inference using OpenVINO"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can use `InferenceOptimizer.trace(..., accelerator='openvino')` API to enable the OpenVINO acceleration for PyTorch inference. It only takes a few lines."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "nbsphinx": "hidden"
   },
   "source": [
    "To apply OpenVINO acceleration, you need to install BigDL-Nano for PyTorch inference first："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx": "hidden"
   },
   "outputs": [],
   "source": [
    "!pip install --pre --upgrade bigdl-nano[pytorch,inference] # install the nightly-built version\n",
    "!source bigdl-nano-init"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nbsphinx": "hidden"
   },
   "source": [
    "> 📝 **Note**\n",
    "> \n",
    "> We recommend to run the commands above, especially `source bigdl-nano-init` before jupyter kernel is started, or some of the optimizations may not take effect."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take an [ResNet-18 model](https://pytorch.org/vision/main/models/generated/torchvision.models.resnet18.html) pretrained on ImageNet dataset as an example. First, we load the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.models import resnet18\n",
    "\n",
    "model_ft = resnet18(pretrained=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To enable OpenVINO acceleration for your PyTorch inference pipeline, **the major change you need to make is to import BigDL-Nano** `InferenceOptimizer`**, and trace your PyTorch model to convert it into an OpenVINO accelerated model for inference**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from bigdl.nano.pytorch import InferenceOptimizer\n",
    "\n",
    "ov_model = InferenceOptimizer.trace(model_ft,\n",
    "                         accelerator=\"openvino\",\n",
    "                         input_sample=torch.rand(1, 3, 224, 224))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 📝 **Note**\n",
    "> \n",
    "> `input_sample` is the parameter for OpenVINO accelerator to know the **shape** of the model input. So both the batch size and the specific values are not important to `input_sample`. If we want our test dataset to consist of images with $224 \\times 224$ pixels, we could use `torch.rand(1, 3, 224, 224)` for `input_sample` here.\n",
    "> \n",
    "> Please refer to [API documentation](https://bigdl.readthedocs.io/en/latest/doc/PythonAPI/Nano/pytorch.html#bigdl.nano.pytorch.InferenceOptimizer.trace) for more information on `InferenceOptimizer.trace`."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You could then do the normal inference steps **under the context manager provided by Nano**, with the model optimized by OpenVINO:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with InferenceOptimizer.get_context(ov_model):\n",
    "    x = torch.rand(2, 3, 224, 224)\n",
    "    # use the optimized model here\n",
    "    y_hat = ov_model(x)\n",
    "    predictions = y_hat.argmax(dim=1)\n",
    "    print(predictions)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 📝 **Note**\n",
    "> \n",
    "> For all Nano optimized models by `InferenceOptimizer.trace`, you need to wrap the inference steps with an automatic context manager `InferenceOptimizer.get_context(model=...)` provided by Nano. You could refer to [here](https://bigdl.readthedocs.io/en/latest/doc/Nano/Howto/Inference/PyTorch/pytorch_context_manager.html) for more detailed usage of the context manager."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 📚 **Related Readings**\n",
    "> \n",
    "> - [How to install BigDL-Nano](https://bigdl.readthedocs.io/en/latest/doc/Nano/Overview/install.html)\n",
    "> - [How to enable automatic context management for PyTorch inference on Nano optimized models](https://bigdl.readthedocs.io/en/latest/doc/Nano/Howto/Inference/PyTorch/pytorch_context_manager.html)\n",
    "> - [How to save and load optimized OpenVINO model](https://bigdl.readthedocs.io/en/latest/doc/Nano/Howto/Inference/PyTorch/pytorch_save_and_load_openvino.html)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nano-pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.7.15 (default, Nov 24 2022, 21:12:53) \n[GCC 11.2.0]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "8a5edab282632443219e051e4ade2d1d5bbc671c781051bf1437897cbdfea0f1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
